package org.dromara.soul.web.request;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.UnpooledByteBufAllocator;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.NettyDataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

/**
* @Description
* @Author hex
* @CreateDate 2020/3/9 16:33
*/
public class CustomServerHttpRequestDecorator extends ServerHttpRequestDecorator {

	private boolean firstSubscribe = true;
	private List<byte[]> listByteArray = new ArrayList<>();

	public CustomServerHttpRequestDecorator(ServerHttpRequest serverHttpRequest) {
		super(serverHttpRequest);
	}

	@Override
	public Flux<DataBuffer> getBody() {
		if (firstSubscribe) {
			return super.getBody().map(dataBuffer -> {
				firstSubscribe = false;
				ByteBuf copy = ((NettyDataBuffer) dataBuffer).getNativeBuffer().copy();
				byte[] dst = new byte[copy.readableBytes()];
				copy.readBytes(dst);
				listByteArray.add(dst);
				copy.release();
				return dataBuffer;
			});
		} else {
			return Flux.fromIterable(listByteArray).map(bytes -> getDataBuffer(bytes));
		}
	}

	private DataBuffer getDataBuffer(byte[] bytes) {
		NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(new UnpooledByteBufAllocator(false));
		return nettyDataBufferFactory.wrap(bytes);
	}
}
